module ZippedArrays

import Base: @propagate_inbounds, size, getindex, setindex

export ZippedArray
"""
    ZippedArray

Zip many arrays of the same size together into a single array.

```jldoctest
julia> using ZippedArrays, Base.Test

julia> array1 = [1 2 3; 4 5 6; 7 8 9]
3×3 Array{Int64,2}:
 1  2  3
 4  5  6
 7  8  9

julia> array2 = [1 4 7; 2 5 8; 3 6 9]
3×3 Array{Int64,2}:
 1  4  7
 2  5  8
 3  6  9

julia> z = @inferred ZippedArray((array1, array2))
3×3 ZippedArrays.ZippedArray{Tuple{Int64,Int64},2,Tuple{Array{Int64,2},Array{Int64,2}}}:
 (1, 1)  (2, 4)  (3, 7)
 (4, 2)  (5, 5)  (6, 8)
 (7, 3)  (8, 6)  (9, 9)

julia> z[2]
(4, 2)

julia> z[2] = (3, 3);

julia> z[2]
(3, 3)
```
"""
struct ZippedArray{T, N, A} <: AbstractArray{T, N}
    arrays::A
end

ZippedArray(arrays::A) where A =
    ZippedArray{Tuple{eltype.(arrays)...}, ndims(first(arrays)), A}(arrays)

@propagate_inbounds getindex_over(vectors, index) =
    map(
        let index = index
            vector -> vector[index...]
        end,
        vectors
    )

@propagate_inbounds setindex_over!(vectors, values, index) =
    map(
        let index = index
            (vector, value) -> vector[index...] = value
        end,
        vectors,
        values
    )

Base.size(e::ZippedArray) = size(first(e.arrays))
Base.getindex(e::ZippedArray{T, N}, index::Vararg{Int, N}) where {T, N} =
    getindex_over(e.arrays, index)
Base.setindex!(e::ZippedArray{T, N}, values, index::Vararg{Int, N}) where {T, N} =
    setindex_over!(e.arrays, values, index)

@propagate_inbounds getindex_over_views(vectors, index) =
    map(
        let index = index
            vector -> @view vector[index...]
        end,
        vectors
    )

export ZippedViews
"""
    ZippedViews

Zip views of many arrays of the same size together into a single array of views.

```jldoctest
julia> using ZippedArrays, Base.Test

julia> array1 = [1 2 3; 4 5 6; 7 8 9]
3×3 Array{Int64,2}:
 1  2  3
 4  5  6
 7  8  9

julia> array2 = [1 4 7; 2 5 8; 3 6 9]
3×3 Array{Int64,2}:
 1  4  7
 2  5  8
 3  6  9

julia> indexes = [ (1, 1:3), (2, 1:3), (3, 1:3) ];

julia> z = @inferred ZippedViews((array1, array2), indexes)
3-element ZippedArrays.ZippedViews{Tuple{Array{Int64,2},Array{Int64,2}},1,Array{Tuple{Int64,UnitRange{Int64}},1}}:
 ([1, 2, 3], [1, 4, 7])
 ([4, 5, 6], [2, 5, 8])
 ([7, 8, 9], [3, 6, 9])

julia> z[2]
([4, 5, 6], [2, 5, 8])

julia> z[2] = ([6, 5, 4], [8, 4, 2]);

julia> z[2]
([6, 5, 4], [8, 4, 2])
```
"""
struct ZippedViews{T, N, L} <: AbstractArray{T, N}
    arrays::T
    locations::L
end

ZippedViews(arrays, locations) =
    ZippedViews{
        typeof(arrays),
        ndims(locations),
        typeof(locations)
    }(arrays, locations)

Base.size(z::ZippedViews) = size(z.locations)
Base.getindex(z::ZippedViews{T, N}, index::Vararg{Int, N}) where {T, N} =
    getindex_over_views(z.arrays, z.locations[index...])
Base.setindex!(z::ZippedViews{T, N}, values, index::Vararg{Int, N}) where {T, N} =
    setindex_over!(z.arrays, values, z.locations[index...])

end
